import argparse
import logging
import os
import torch

from unet import UNet_onnx,UNet


def get_args():
    parser = argparse.ArgumentParser(description='Predict masks from input images')
    parser.add_argument('--model', '-m',
                        default='checkpoints/5_contrast+sampling+attention/best_checkpoint_epoch162_train0.0023_val0.9939.pth',
                        metavar='FILE',
                        help='Specify the file in which the model is stored')
    parser.add_argument('--bilinear', action='store_true', default=False, help='Use bilinear upsampling')
    parser.add_argument('--classes', '-c', type=int, default=2, help='Number of classes')
    parser.add_argument('--is_edges', '-ed', type=bool, default=True,
                        help='Is edge detection used for image preprocessing')

    return parser.parse_args()


if __name__ == '__main__':
    args = get_args()
    onnx_path='best_checkpoint_epoch162_train0.0023_val0.9939.onnx'

    if args.is_edges:
        channels = 4
    else:
        channels = 3

    net = UNet_onnx(n_channels=channels, n_classes=args.classes, bilinear=args.bilinear)

    device=torch.device("cpu")
    state_dict = torch.load(args.model,map_location=device)
    mask_values = state_dict.pop('mask_values', [0, 1])
    net.load_state_dict(state_dict)

    logging.info('Model loaded!')
    # 定义输入数据
    dummy_input = torch.randn(1, channels, 320, 480)
    # 使用预训练模型进行前向传播
    # output = net(dummy_input)

    # 导出模型为ONNX格式
    torch.onnx.export(net, dummy_input, onnx_path)
